#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
火灾检测API服务
基于realtime_fire_detection.py的火灾检测系统，提供REST API接口
"""

import os
import io
import base64
import json
import time
import logging
import socket
import pickle
import struct
import threading
import queue
from datetime import datetime
from typing import Dict, Any, Optional
import argparse

import cv2
import numpy as np
import requests
from PIL import Image
from flask import Flask, request, jsonify, Response
from flask_socketio import SocketIO, emit
from werkzeug.utils import secure_filename

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('fire_detection_api.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class FireDetectionSystem:
    """基于Ollama Gemma3模型的火灾检测系统"""
    
    def __init__(self, ollama_url: str = "http://localhost:11434", model_name: str = "gemma3:latest", debug: bool = False):
        self.ollama_url = ollama_url
        self.model_name = model_name
        self.api_url = f"{ollama_url}/api/generate"
        self.debug = debug
        
    def check_ollama_connection(self) -> bool:
        """检查Ollama服务是否可用"""
        try:
            response = requests.get(f"{self.ollama_url}/api/tags", timeout=5)
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False
    
    def detect_fire_from_image(self, image_data: np.ndarray) -> Dict[str, Any]:
        """
        从图像数据检测火灾
        
        Args:
            image_data: OpenCV图像数据
            
        Returns:
            检测结果字典
        """
        try:
            # 将OpenCV帧转换为PIL图像
            if len(image_data.shape) == 3:
                frame_rgb = cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)
            else:
                frame_rgb = image_data
            pil_image = Image.fromarray(frame_rgb)
            
            # 调整图片大小
            pil_image = pil_image.resize((512, 512), Image.Resampling.LANCZOS)
            
            # 转换为base64
            buffer = io.BytesIO()
            pil_image.save(buffer, format='JPEG', quality=85)
            image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            # 创建提示词
            prompt = "请仔细观察这张图片，判断图中是否发生了火灾。请直接回答：是 或 否。"
            
            # 准备请求数据
            payload = {
                "model": self.model_name,
                "prompt": prompt,
                "images": [image_base64],
                "stream": False
            }
            
            # 发送请求
            response = requests.post(
                self.api_url,
                json=payload,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                response_text = result.get('response', '')
                
                # 判断逻辑
                if '是' in response_text and '否' not in response_text:
                    has_fire = True
                    confidence = 0.9
                elif '否' in response_text and '是' not in response_text:
                    has_fire = False
                    confidence = 0.9
                elif '是' in response_text and '否' in response_text:
                    has_fire = True
                    confidence = 0.7
                else:
                    has_fire = False
                    confidence = 0.7
                
                return {
                    "has_fire": has_fire,
                    "confidence": confidence,
                    "raw_response": response_text,
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "status": "success"
                }
            else:
                return {
                    "error": f"Ollama API请求失败: {response.status_code}",
                    "has_fire": False,
                    "confidence": 0.0,
                    "status": "error"
                }
                
        except Exception as e:
            logger.error(f"火灾检测失败: {str(e)}")
            return {
                "error": f"检测过程中发生错误: {str(e)}",
                "has_fire": False,
                "confidence": 0.0,
                "status": "error"
            }

class VideoStreamHandler:
    """视频流处理器"""
    
    def __init__(self, fire_detector, socketio=None, detection_interval=3.0):
        self.fire_detector = fire_detector
        self.socketio = socketio
        self.detection_interval = detection_interval
        self.frame_queue = queue.Queue(maxsize=10)
        self.running = False
        self.detection_thread = None
        self.last_detection_time = 0
        self.current_frame = None
        self.frame_lock = threading.Lock()
        
        # 视频帧发送控制
        self.last_frame_send_time = 0
        self.frame_send_interval = 1.0 / 30.0  # 30fps
        
    def start_detection(self):
        """启动检测线程"""
        if not self.running:
            self.running = True
            self.detection_thread = threading.Thread(target=self.detection_worker, daemon=True)
            self.detection_thread.start()
            logger.info("视频流检测线程已启动")
    
    def stop_detection(self):
        """停止检测线程"""
        self.running = False
        if self.detection_thread:
            self.detection_thread.join(timeout=2)
        logger.info("视频流检测线程已停止")
    
    def add_frame(self, frame):
        """添加帧到检测队列"""
        try:
            # 存储当前帧用于WebSocket广播
            with self.frame_lock:
                self.current_frame = frame.copy()
            
            # 将帧添加到检测队列
            if not self.frame_queue.full():
                self.frame_queue.put_nowait((frame, time.time()))
            else:
                # 队列满时丢弃最旧的帧
                try:
                    self.frame_queue.get_nowait()
                    self.frame_queue.put_nowait((frame, time.time()))
                except queue.Empty:
                    pass
            
            # 通过WebSocket广播当前帧（带帧率控制）
            if self.socketio:
                current_time = time.time()
                if current_time - self.last_frame_send_time >= self.frame_send_interval:
                    self.broadcast_frame(frame)
                    self.last_frame_send_time = current_time
                
        except Exception as e:
            logger.error(f"添加帧失败: {str(e)}")
    
    def broadcast_frame(self, frame):
        """通过WebSocket广播视频帧"""
        try:
            # 将帧编码为JPEG
            _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
            frame_base64 = base64.b64encode(buffer).decode('utf-8')
            
            # 通过WebSocket发送帧数据
            self.socketio.emit('video_frame', {
                'frame': frame_base64,
                'timestamp': time.time()
            })
        except Exception as e:
            logger.error(f"广播帧失败: {str(e)}")
    
    def get_current_frame(self):
        """获取当前帧"""
        with self.frame_lock:
            return self.current_frame.copy() if self.current_frame is not None else None
    
    def detection_worker(self):
        """检测工作线程"""
        while self.running:
            try:
                frame_data = self.frame_queue.get(timeout=1.0)
                frame, timestamp = frame_data
                
                # 检查检测间隔
                current_time = time.time()
                if current_time - self.last_detection_time >= self.detection_interval:
                    # 执行火灾检测
                    detection_result = self.fire_detector.detect_fire_from_image(frame)
                    
                    # 记录检测结果
                    has_fire = detection_result.get('has_fire', False)
                    confidence = detection_result.get('confidence', 0.0)
                    
                    if has_fire:
                        logger.warning(f"🔥 检测到火灾迹象 - 置信度: {confidence:.2f}")
                    else:
                        logger.info(f"✅ 安全 - 置信度: {confidence:.2f}")
                    
                    # 通过WebSocket广播检测结果
                    if self.socketio:
                        self.socketio.emit('detection_result', {
                            'has_fire': has_fire,
                            'confidence': confidence,
                            'raw_response': detection_result.get('raw_response', ''),
                            'timestamp': detection_result.get('timestamp', ''),
                            'status': 'success'
                        })
                    
                    self.last_detection_time = current_time
                
                self.frame_queue.task_done()
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"检测工作线程错误: {str(e)}")

class FireDetectionAPI:
    """火灾检测API服务"""
    
    def __init__(self, ollama_url: str = "http://localhost:11434", model_name: str = "gemma3:latest", 
                 video_port: int = 5000):
        self.app = Flask(__name__)
        self.app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
        self.app.config['SECRET_KEY'] = 'fire_detection_secret_key'
        
        # 初始化SocketIO
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")
        
        # 初始化火灾检测系统
        self.fire_detector = FireDetectionSystem(ollama_url, model_name)
        
        # 初始化视频流处理器
        self.video_handler = VideoStreamHandler(self.fire_detector, self.socketio)
        
        # 视频流服务器设置
        self.video_port = video_port
        self.video_server = None
        self.video_running = False
        
        # 检查Ollama连接
        self.check_ollama_connection()
        
        # 设置路由
        self.setup_routes()
        
        # 设置WebSocket事件
        self.setup_socketio_events()
        
        # 创建上传目录
        self.upload_dir = "uploads"
        if not os.path.exists(self.upload_dir):
            os.makedirs(self.upload_dir)
            logger.info(f"创建上传目录: {self.upload_dir}")
    
    def check_ollama_connection(self):
        """检查Ollama连接"""
        if self.fire_detector.check_ollama_connection():
            logger.info("✅ Ollama服务连接正常")
        else:
            logger.error("❌ Ollama服务不可用，请确保Ollama正在运行")
            logger.error("安装命令: ollama pull gemma3:latest")
    
    def start_video_server(self):
        """启动视频流服务器"""
        if not self.video_running:
            self.video_running = True
            self.video_server = threading.Thread(target=self.video_server_worker, daemon=True)
            self.video_server.start()
            self.video_handler.start_detection()
            logger.info(f"视频流服务器启动在端口 {self.video_port}")
    
    def stop_video_server(self):
        """停止视频流服务器"""
        self.video_running = False
        self.video_handler.stop_detection()
        if self.video_server:
            self.video_server.join(timeout=2)
        logger.info("视频流服务器已停止")
    
    def video_server_worker(self):
        """视频流服务器工作线程"""
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        
        try:
            server_socket.bind(('0.0.0.0', self.video_port))
            server_socket.listen(5)
            server_socket.settimeout(1.0)
            
            while self.video_running:
                try:
                    client_socket, addr = server_socket.accept()
                    logger.info(f"视频客户端连接: {addr}")
                    threading.Thread(target=self.handle_video_client, args=(client_socket,), daemon=True).start()
                except socket.timeout:
                    continue
                except Exception as e:
                    if self.video_running:
                        logger.error(f"接受视频连接时出错: {e}")
                    break
        except Exception as e:
            logger.error(f"视频服务器启动失败: {e}")
        finally:
            server_socket.close()
    
    def handle_video_client(self, client_socket):
        """处理视频客户端连接"""
        payload_size = struct.calcsize("L")
        data = b""
        
        try:
            while self.video_running:
                client_socket.settimeout(5.0)
                
                try:
                    # 接收消息大小
                    while len(data) < payload_size:
                        chunk = client_socket.recv(4096)
                        if not chunk:
                            logger.info("视频客户端断开连接")
                            break
                        data += chunk
                    
                    # 检查是否收到足够的数据
                    if len(data) < payload_size:
                        logger.info("视频数据不完整，客户端可能断开连接")
                        break
                        
                    packed_msg_size = data[:payload_size]
                    data = data[payload_size:]
                    msg_size = struct.unpack("L", packed_msg_size)[0]
                    
                    # 处理心跳包
                    if msg_size == 0:
                        continue
                    
                    # 接收帧数据
                    while len(data) < msg_size:
                        chunk = client_socket.recv(4096)
                        if not chunk:
                            logger.info("视频客户端断开连接")
                            break
                        data += chunk
                    
                    # 检查是否收到完整的帧数据
                    if len(data) < msg_size:
                        logger.info("视频帧数据不完整，客户端可能断开连接")
                        break
                        
                    frame_data = data[:msg_size]
                    data = data[msg_size:]
                    
                    # 反序列化帧数据
                    try:
                        buffer = pickle.loads(frame_data)
                        frame = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
                        
                        if frame is not None:
                            # 将帧添加到检测队列
                            self.video_handler.add_frame(frame)
                            
                    except Exception as e:
                        logger.warning(f"视频帧数据解析失败: {e}")
                        continue
                        
                except socket.timeout:
                    continue
                except ConnectionError as e:
                    logger.info(f"视频客户端连接断开: {e}")
                    break
                    
        except Exception as e:
            logger.error(f"处理视频客户端错误: {e}")
        finally:
            client_socket.close()
            logger.info("视频客户端连接已关闭")
    
    def setup_routes(self):
        """设置API路由"""
        
        @self.app.route('/health', methods=['GET'])
        def health_check():
            """健康检查接口"""
            ollama_status = self.fire_detector.check_ollama_connection()
            return jsonify({
                "status": "healthy" if ollama_status else "unhealthy",
                "ollama_connected": ollama_status,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
        
        @self.app.route('/detect', methods=['POST'])
        def detect_fire():
            """火灾检测接口"""
            try:
                # 检查是否有文件上传
                if 'image' not in request.files:
                    return jsonify({
                        "error": "没有上传图片文件",
                        "status": "error"
                    }), 400
                
                file = request.files['image']
                if file.filename == '':
                    return jsonify({
                        "error": "没有选择文件",
                        "status": "error"
                    }), 400
                
                # 检查文件类型
                if not self.is_allowed_file(file.filename):
                    return jsonify({
                        "error": "不支持的文件类型，请上传图片文件",
                        "status": "error"
                    }), 400
                
                # 读取图片数据
                file_data = file.read()
                nparr = np.frombuffer(file_data, np.uint8)
                image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                
                if image is None:
                    return jsonify({
                        "error": "无法解析图片文件",
                        "status": "error"
                    }), 400
                
                # 执行火灾检测
                result = self.fire_detector.detect_fire_from_image(image)
                
                # 记录检测结果
                logger.info(f"火灾检测结果: {result}")
                
                return jsonify(result)
                
            except Exception as e:
                logger.error(f"检测接口错误: {str(e)}")
                return jsonify({
                    "error": f"服务器内部错误: {str(e)}",
                    "status": "error"
                }), 500
        
        @self.app.route('/detect_base64', methods=['POST'])
        def detect_fire_base64():
            """基于Base64的火灾检测接口"""
            try:
                data = request.get_json()
                if not data or 'image' not in data:
                    return jsonify({
                        "error": "请求中缺少image字段",
                        "status": "error"
                    }), 400
                
                # 解码Base64图片
                try:
                    image_data = base64.b64decode(data['image'])
                    nparr = np.frombuffer(image_data, np.uint8)
                    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                except Exception as e:
                    return jsonify({
                        "error": f"Base64解码失败: {str(e)}",
                        "status": "error"
                    }), 400
                
                if image is None:
                    return jsonify({
                        "error": "无法解析图片数据",
                        "status": "error"
                    }), 400
                
                # 执行火灾检测
                result = self.fire_detector.detect_fire_from_image(image)
                
                # 记录检测结果
                logger.info(f"火灾检测结果: {result}")
                
                return jsonify(result)
                
            except Exception as e:
                logger.error(f"Base64检测接口错误: {str(e)}")
                return jsonify({
                    "error": f"服务器内部错误: {str(e)}",
                    "status": "error"
                }), 500
        
        @self.app.route('/detect_url', methods=['POST'])
        def detect_fire_url():
            """基于URL的火灾检测接口"""
            try:
                data = request.get_json()
                if not data or 'url' not in data:
                    return jsonify({
                        "error": "请求中缺少url字段",
                        "status": "error"
                    }), 400
                
                # 从URL下载图片
                try:
                    response = requests.get(data['url'], timeout=10)
                    response.raise_for_status()
                    
                    nparr = np.frombuffer(response.content, np.uint8)
                    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                except Exception as e:
                    return jsonify({
                        "error": f"下载图片失败: {str(e)}",
                        "status": "error"
                    }), 400
                
                if image is None:
                    return jsonify({
                        "error": "无法解析图片数据",
                        "status": "error"
                    }), 400
                
                # 执行火灾检测
                result = self.fire_detector.detect_fire_from_image(image)
                
                # 记录检测结果
                logger.info(f"火灾检测结果: {result}")
                
                return jsonify(result)
                
            except Exception as e:
                logger.error(f"URL检测接口错误: {str(e)}")
                return jsonify({
                    "error": f"服务器内部错误: {str(e)}",
                    "status": "error"
                }), 500
        
        @self.app.route('/info', methods=['GET'])
        def get_info():
            """获取服务信息"""
            return jsonify({
                "service": "火灾检测API服务",
                "version": "1.0.0",
                "ollama_url": self.fire_detector.ollama_url,
                "model_name": self.fire_detector.model_name,
                "video_port": self.video_port,
                "video_running": self.video_running,
                "endpoints": {
                    "health": "/health - 健康检查",
                    "detect": "/detect - 文件上传检测",
                    "detect_base64": "/detect_base64 - Base64图片检测",
                    "detect_url": "/detect_url - URL图片检测",
                    "video_start": "/video/start - 启动视频流服务",
                    "video_stop": "/video/stop - 停止视频流服务",
                    "video_status": "/video/status - 视频流状态",
                    "info": "/info - 服务信息"
                },
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
        
        @self.app.route('/video/start', methods=['POST'])
        def start_video():
            """启动视频流服务"""
            try:
                if not self.video_running:
                    self.start_video_server()
                    return jsonify({
                        "status": "success",
                        "message": "视频流服务已启动",
                        "video_port": self.video_port,
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
                else:
                    return jsonify({
                        "status": "info",
                        "message": "视频流服务已在运行",
                        "video_port": self.video_port,
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
            except Exception as e:
                logger.error(f"启动视频流服务失败: {e}")
                return jsonify({
                    "status": "error",
                    "message": f"启动视频流服务失败: {str(e)}",
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }), 500
        
        @self.app.route('/video/stop', methods=['POST'])
        def stop_video():
            """停止视频流服务"""
            try:
                if self.video_running:
                    self.stop_video_server()
                    return jsonify({
                        "status": "success",
                        "message": "视频流服务已停止",
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
                else:
                    return jsonify({
                        "status": "info",
                        "message": "视频流服务未运行",
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
            except Exception as e:
                logger.error(f"停止视频流服务失败: {e}")
                return jsonify({
                    "status": "error",
                    "message": f"停止视频流服务失败: {str(e)}",
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }), 500
        
        @self.app.route('/video/status', methods=['GET'])
        def video_status():
            """获取视频流状态"""
            return jsonify({
                "video_running": self.video_running,
                "video_port": self.video_port,
                "detection_running": self.video_handler.running,
                "queue_size": self.video_handler.frame_queue.qsize(),
                "last_detection_time": self.video_handler.last_detection_time,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
        
        @self.app.route('/video/frame', methods=['GET'])
        def get_current_frame():
            """获取当前视频帧"""
            try:
                current_frame = self.video_handler.get_current_frame()
                if current_frame is not None:
                    # 将帧编码为JPEG
                    _, buffer = cv2.imencode('.jpg', current_frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                    frame_base64 = base64.b64encode(buffer).decode('utf-8')
                    
                    return jsonify({
                        "success": True,
                        "frame": frame_base64,
                        "timestamp": time.time()
                    })
                else:
                    return jsonify({
                        "success": False,
                        "error": "没有可用的视频帧"
                    })
            except Exception as e:
                return jsonify({
                    "success": False,
                    "error": str(e)
                })
    
    def is_allowed_file(self, filename: str) -> bool:
        """检查文件类型是否允许"""
        ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp'}
        return '.' in filename and \
               filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
    
    def setup_socketio_events(self):
        """设置WebSocket事件处理"""
        
        @self.socketio.on('connect')
        def handle_connect():
            logger.info('客户端已连接')
            emit('connected', {'message': '已连接到火灾检测服务'})
        
        @self.socketio.on('disconnect')
        def handle_disconnect():
            logger.info('客户端已断开连接')
        
        @self.socketio.on('request_video_stream')
        def handle_video_stream_request():
            """处理视频流请求"""
            if self.video_handler.running:
                emit('video_stream_status', {'status': 'active'})
            else:
                emit('video_stream_status', {'status': 'inactive'})
        
        @self.socketio.on('start_detection')
        def handle_start_detection():
            """处理开始检测请求"""
            if not self.video_handler.running:
                self.video_handler.start_detection()
                emit('detection_status', {'status': 'started'})
            else:
                emit('detection_status', {'status': 'already_running'})
        
        @self.socketio.on('stop_detection')
        def handle_stop_detection():
            """处理停止检测请求"""
            if self.video_handler.running:
                self.video_handler.stop_detection()
                emit('detection_status', {'status': 'stopped'})
            else:
                emit('detection_status', {'status': 'not_running'})
    
    def run(self, host='0.0.0.0', port=8080, debug=False, auto_start_video=True):
        """启动API服务"""
        logger.info(f"🔥 火灾检测API服务启动在 {host}:{port}")
        logger.info("可用的API端点:")
        logger.info("  GET  /health - 健康检查")
        logger.info("  POST /detect - 文件上传检测")
        logger.info("  POST /detect_base64 - Base64图片检测")
        logger.info("  POST /detect_url - URL图片检测")
        logger.info("  POST /video/start - 启动视频流服务")
        logger.info("  POST /video/stop - 停止视频流服务")
        logger.info("  GET  /video/status - 视频流状态")
        logger.info("  GET  /info - 服务信息")
        
        # 自动启动视频流服务
        if auto_start_video:
            logger.info("自动启动视频流服务...")
            self.start_video_server()
        
        try:
            self.socketio.run(self.app, host=host, port=port, debug=debug)
        except KeyboardInterrupt:
            logger.info("收到中断信号，正在停止服务...")
            self.stop_video_server()
            logger.info("服务已停止")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='火灾检测API服务')
    parser.add_argument('--host', default='0.0.0.0', help='服务器绑定地址')
    parser.add_argument('--port', type=int, default=8080, help='API监听端口号')
    parser.add_argument('--video-port', type=int, default=5000, help='视频流端口号')
    parser.add_argument('--ollama-url', default='http://localhost:11434', help='Ollama服务地址')
    parser.add_argument('--model', default='gemma3:latest', help='Ollama模型名称')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    parser.add_argument('--no-auto-video', action='store_true', help='不自动启动视频流服务')
    
    args = parser.parse_args()
    
    print("🔥 火灾检测API服务")
    print("=" * 50)
    print(f"API服务器地址: {args.host}:{args.port}")
    print(f"视频流端口: {args.video_port}")
    print(f"Ollama地址: {args.ollama_url}")
    print(f"模型名称: {args.model}")
    print(f"调试模式: {'启用' if args.debug else '禁用'}")
    print(f"自动启动视频流: {'禁用' if args.no_auto_video else '启用'}")
    print("=" * 50)
    
    # 创建并启动API服务
    api_service = FireDetectionAPI(args.ollama_url, args.model, args.video_port)
    api_service.run(
        host=args.host, 
        port=args.port, 
        debug=args.debug,
        auto_start_video=not args.no_auto_video
    )

if __name__ == "__main__":
    main()
